# fl_client.py
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from isomorphic_noise_generator import IsomorphicNoiseGenerator
from quantum_audit_system import QuantumAuditSystem
from dynamic_noise_optimizer import DynamicNoiseOptimizer
from noise_revocation_engine import NoiseRevocationEngine

class FLClient:
    """
    Federated learning client with noise revocation capabilities
    
    Args:
        client_id (int): Unique client identifier
        data (Dataset): Local training dataset
        model (nn.Module): Local model instance
        seed (bytes): Shared secret seed
    """
    
    def __init__(self, client_id, data, model, seed):
        self.client_id = client_id
        self.data_loader = DataLoader(data, batch_size=32, shuffle=True)
        self.model = model
        self.noise_gen = IsomorphicNoiseGenerator(seed, self._count_parameters())
        self.quantum_audit = QuantumAuditSystem()
        self.noise_optimizer = DynamicNoiseOptimizer()
    
    def _count_parameters(self):
        """Count trainable parameters in model"""
        return sum(p.numel() for p in self.model.parameters() if p.requires_grad)
    
    def local_train(self, round_num, total_rounds):
        """
        Perform local training with noise injection
        
        Args:
            round_num (int): Current training round
            total_rounds (int): Total number of training rounds
            
        Returns:
            tuple: (noisy_params, quantum_tag)
        """
        optimizer = optim.SGD(self.model.parameters(), lr=0.01)
        criterion = nn.CrossEntropyLoss()
        
        # Update dynamic parameters
        self.quantum_audit.update_entropy(round_num, total_rounds)
        self.noise_optimizer.update_clustering_threshold(round_num, total_rounds)
        
        # Generate quantum audit tag
        quantum_tag = self.quantum_audit.generate_quantum_tag(self.client_id, round_num)
        
        # Train model locally
        for inputs, labels in self.data_loader:
            optimizer.zero_grad()
            outputs = self.model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
        
        # Generate and apply isomorphic noise
        noise_vector = self.noise_gen.generate_noise(round_num)
        noisy_params = NoiseRevocationEngine.apply_noise(
            list(self.model.parameters()), noise_vector
        )
        
        return noisy_params, quantum_tag
